// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <atomic>
#include <cstdint>
#include <map>
#include <memory>
#include <unordered_set>

#include <pollux/common/memory/memory_allocator.h>

namespace kumo::pollux::memory {
    class MmapArena {
    public:
        /// Single MmapArena capacity is determined by mmap_arena_capacity_ratio ratio
        /// but the capacity should be at least kMinCapacityBytes if the calculated
        /// capacity from the ratio is too small to serve large allocations.
        static constexpr int64_t kMinCapacityBytes = 128 * 1024 * 1024; // 128M

        /// MmapArena capacity should be multiple of kMinGrainSizeBytes.
        static constexpr uint64_t kMinGrainSizeBytes = 1024 * 1024; // 1M

        explicit MmapArena(size_t capacityBytes);

        ~MmapArena();

        void *allocate(uint64_t bytes);

        void free(void *address, uint64_t bytes);

        void *address() const {
            return reinterpret_cast<void *>(address_);
        }

        uint64_t byteSize() const {
            return byteSize_;
        }

        const std::map<uintptr_t, uint64_t> &freeList() const {
            return freeList_;
        }

        const std::map<uint64_t, std::unordered_set<uintptr_t> > &freeLookup() const {
            return freeLookup_;
        }

        uint64_t freeBytes() {
            return freeBytes_;
        }

        bool empty() {
            return freeBytes_ == byteSize_;
        }

        /// Checks internal consistency of this MmapArena. Returns true if OK. May
        /// return false if there are concurrent allocations and frees during the
        /// consistency check. This is a false positive but not dangerous. This is for
        /// test only
        bool checkConsistency() const;

        /// translate lookup table to a string for debugging purpose only.
        std::string freeLookupStr() {
            std::stringstream lookupStr;
            for (auto itr = freeLookup_.begin(); itr != freeLookup_.end(); ++itr) {
                lookupStr << "\n{" << itr->first << "->[";
                for (auto itrInner = itr->second.begin(); itrInner != itr->second.end();
                     itrInner++) {
                    lookupStr << *itrInner << ", ";
                }
                lookupStr << "]}\n";
            }
            return lookupStr.str();
        }

        std::string toString() const;

    private:
        // Rounds up size to the next power of 2.
        static uint64_t roundBytes(uint64_t bytes);

        std::map<uintptr_t, uint64_t>::iterator addFreeBlock(
            uintptr_t addr,
            uint64_t bytes);

        void removeFromLookup(uintptr_t addr, uint64_t bytes);

        void removeFreeBlock(uintptr_t addr, uint64_t bytes);

        void removeFreeBlock(std::map<uintptr_t, uint64_t>::iterator &itr);

        // Total capacity size of this arena.
        const uint64_t byteSize_;

        // Starting address of this arena.
        uint8_t *address_;

        std::atomic<uint64_t> freeBytes_;

        // A sorted list with each entry mapping from free block address to size of
        // the free block
        std::map<uintptr_t, uint64_t> freeList_;

        // A sorted look-up structure that stores the block size as key and a set of
        // addresses of that size as value.
        std::map<uint64_t, std::unordered_set<uintptr_t> > freeLookup_;
    };

    /// A class that manages a set of MmapArenas. It is able to adapt itself by
    /// growing the number of its managed MmapArena's when extreme memory
    /// fragmentation happens.
    class ManagedMmapArenas {
    public:
        explicit ManagedMmapArenas(uint64_t singleArenaCapacity);

        void *allocate(uint64_t bytes);

        void free(void *address, uint64_t bytes);

        const std::map<uintptr_t, std::shared_ptr<MmapArena> > &arenas() const {
            return arenas_;
        }

    private:
        // Capacity in bytes for a single MmapArena managed by this.
        const uint64_t singleArenaCapacity_;

        // A sorted list of MmapArena by its initial address
        std::map<uintptr_t, std::shared_ptr<MmapArena> > arenas_;

        // All allocations should come from this MmapArena. When it is no longer able
        // to handle allocations it will be updated to a newly created MmapArena.
        std::shared_ptr<MmapArena> currentArena_;
    };
} // namespace kumo::pollux::memory
